package main

import (
	"log"
	"net"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"github.com/nmcclain/ldap"
)

// 用户信息结构
type User struct {
	CN       string
	DeptName string
	Username string
	Password string
	Email    string
	Phone    string
}

// 部门信息结构
type Department struct {
	P_OU        string
	OU          string
	Name        string
	Description string
}

// LDAP目录数据
var (
	domainDN string = "dc=example,dc=com"
	users           = []User{
		{
			CN:       "admin",
			Username: "admin",
			Password: "123456",
			Email:    "ad@demo.com",
			Phone:    "13910003346",
		},
		{
			CN:       "user1-1",
			DeptName: "dept1-1",
			Username: "user1-1",
			Password: "123456"},
		{
			CN:       "user1-2",
			DeptName: "dept1-1",
			Username: "user1-2",
			Password: "123456"},
		{
			CN:       "user1-3",
			DeptName: "dept1-2",
			Username: "user1-3",
			Password: "123456"},
		{
			CN:       "user1-4",
			DeptName: "dept1-2",
			Username: "user1-4",
			Password: "123456"},
		{
			CN:       "user2-1",
			DeptName: "dept2-1",
			Username: "user2-1",
			Password: "123456"},
		{
			CN:       "user2-2",
			DeptName: "dept2-1",
			Username: "user2-2",
			Password: "123456"},
		{
			CN:       "user2-4",
			DeptName: "dept2-2",
			Username: "user2-3",
			Password: "123456"},
	}

	departments = []Department{
		{OU: "dept1", P_OU: "", Name: "dept1", Description: "部门1"},
		{OU: "dept1-1", P_OU: "dept1", Name: "dept1-1", Description: "部门1-1"},
		{OU: "dept1-2", P_OU: "dept1", Name: "dept1-2", Description: "部门1-2"},
		{OU: "dept2", P_OU: "", Name: "dept2", Description: "部门2"},
		{OU: "dept2-1", P_OU: "dept2", Name: "dept2-1", Description: "部门2-1"},
		{OU: "dept2-2", P_OU: "dept2", Name: "dept2-2", Description: "部门2-2"},
	}
)

type simpleBindHandler struct{}

func getUser(username string) (User, bool) {
	for _, user := range users {
		if user.CN == username {
			return user, true
		}
	}
	return User{}, false
}

func getDepartment(name string) (Department, bool) {
	for _, department := range departments {
		if department.OU == name {
			return department, true
		}
	}
	return Department{}, false
}

func getDepartments(p_ou string) []Department {
	var result []Department
	for _, department := range departments {
		if department.P_OU == p_ou {
			result = append(result, department)
		}
	}
	return result
}

func getUsers(ou string) []User {
	var result []User
	for _, user := range users {
		if user.DeptName == ou {
			result = append(result, user)
		}
	}
	return result
}

func (h *simpleBindHandler) Bind(bindDN, bindSimplePw string, conn net.Conn) (ldap.LDAPResultCode, error) {
	log.Printf("Bind request: dn:%s password:%s", bindDN, bindSimplePw)

	// 解析DN获取用户名
	parts := strings.Split(bindDN, ",")
	if len(parts) == 0 {
		return ldap.LDAPResultInvalidCredentials, nil
	}

	// 提取CN部分
	cnPart := strings.Split(parts[0], "=")
	if len(cnPart) != 2 || cnPart[0] != "cn" {
		return ldap.LDAPResultInvalidCredentials, nil
	}

	username := cnPart[1]

	// 验证用户是否存在及密码是否正确
	user, exists := getUser(username)
	if !exists || user.Password != bindSimplePw {
		return ldap.LDAPResultInvalidCredentials, nil
	}

	return ldap.LDAPResultSuccess, nil
}

type simpleSearchHandler struct{}

func (h *simpleSearchHandler) Search(bindDN string, req ldap.SearchRequest, conn net.Conn) (ldap.ServerSearchResult, error) {
	log.Printf("Search request: %s, Scope: %d", req.BaseDN, req.Scope)

	// 获取当前登录用户
	currentUser := ""
	if strings.HasPrefix(bindDN, "cn=") {
		currentUser = strings.Split(strings.Split(bindDN, ",")[0], "=")[1]
	}

	// 根据Scope判断是获取属性信息还是下级节点信息
	// req.Scope: 0=BaseObject, 1=SingleLevel, 2=WholeSubtree
	isAttributeRequest := req.Scope == 0 // BaseObject scope表示获取当前节点属性

	var resultEntries []*ldap.Entry
	var err error

	switch {
	case req.BaseDN == domainDN:
		if isAttributeRequest {
			resultEntries = h.getRootAttributes(req.BaseDN)
		} else {
			resultEntries, err = h.searchDepartment("", currentUser)
		}
	case strings.HasPrefix(req.BaseDN, "ou="):
		ouName := strings.Split(strings.Split(req.BaseDN, ",")[0], "=")[1]
		if isAttributeRequest {
			resultEntries = h.getDepartmentAttributes(req.BaseDN, ouName)
		} else {
			resultEntries, err = h.searchDepartment(ouName, currentUser)
		}
	case strings.HasPrefix(req.BaseDN, "cn="):
		cnName := strings.Split(strings.Split(req.BaseDN, ",")[0], "=")[1]
		if isAttributeRequest {
			resultEntries = h.getUserAttributes(req.BaseDN, cnName)
		}
		// 用户节点没有下级节点，所以不需要处理非属性请求
	}

	if err != nil {
		return ldap.ServerSearchResult{}, err
	}

	return ldap.ServerSearchResult{
		Entries: resultEntries,
	}, nil
}

func (h *simpleSearchHandler) getRootAttributes(baseDN string) []*ldap.Entry {
	return []*ldap.Entry{{
		DN: baseDN,
		Attributes: []*ldap.EntryAttribute{
			{Name: "dc", Values: []string{"example"}},
			{Name: "objectClass", Values: []string{"domain"}},
		},
	}}
}

func (h *simpleSearchHandler) getDepartmentAttributes(baseDN, ouName string) []*ldap.Entry {
	dept, exists := getDepartment(ouName)
	if !exists {
		return nil
	}

	return []*ldap.Entry{{
		DN: baseDN,
		Attributes: []*ldap.EntryAttribute{
			{Name: "ou", Values: []string{dept.OU}},
			{Name: "objectClass", Values: []string{"organizationalUnit"}},
			{Name: "description", Values: []string{dept.Description}},
		},
	}}
}

func (h *simpleSearchHandler) searchDepartment(ouName string, currentUser string) ([]*ldap.Entry, error) {
	var resultEntries []*ldap.Entry

	// 如果是admin用户，返回所有数据
	if currentUser == "admin" {
		// 获取当前部门下的子部门和用户
		depts := getDepartments(ouName)
		for _, dept := range depts {
			deptDN := "ou=" + dept.OU + ",ou=" + ouName + "," + domainDN
			resultEntries = append(resultEntries, &ldap.Entry{
				DN: deptDN,
				Attributes: []*ldap.EntryAttribute{
					{Name: "ou", Values: []string{dept.OU}},
					{Name: "objectClass", Values: []string{"organizationalUnit"}},
					{Name: "description", Values: []string{dept.Description}},
				},
			})
		}

		users := getUsers(ouName)
		for _, user := range users {
			userDN := "cn=" + user.CN + ",ou=" + ouName + "," + domainDN
			resultEntries = append(resultEntries, &ldap.Entry{
				DN: userDN,
				Attributes: []*ldap.EntryAttribute{
					{Name: "cn", Values: []string{user.CN}},
					{Name: "objectClass", Values: []string{"organizationalPerson"}},
					{Name: "mail", Values: []string{user.Email}},
				},
			})
		}
	} else {
		// 普通用户只能查看自己部门的数据
		user, exists := getUser(currentUser)
		if !exists {
			return nil, nil
		}

		// 如果请求的是用户所在部门的数据
		if ouName == user.DeptName {
			// 获取当前部门下的子部门和用户
			depts := getDepartments(ouName)
			for _, dept := range depts {
				deptDN := "ou=" + dept.OU + ",ou=" + ouName + "," + domainDN
				resultEntries = append(resultEntries, &ldap.Entry{
					DN: deptDN,
					Attributes: []*ldap.EntryAttribute{
						{Name: "ou", Values: []string{dept.OU}},
						{Name: "objectClass", Values: []string{"organizationalUnit"}},
						{Name: "description", Values: []string{dept.Description}},
					},
				})
			}

			users := getUsers(ouName)
			for _, user := range users {
				userDN := "cn=" + user.CN + ",ou=" + ouName + "," + domainDN
				resultEntries = append(resultEntries, &ldap.Entry{
					DN: userDN,
					Attributes: []*ldap.EntryAttribute{
						{Name: "cn", Values: []string{user.CN}},
						{Name: "objectClass", Values: []string{"organizationalPerson"}},
						{Name: "mail", Values: []string{user.Email}},
					},
				})
			}
		}
	}

	return resultEntries, nil
}

func (h *simpleSearchHandler) getUserAttributes(baseDN, cnName string) []*ldap.Entry {
	user, exists := getUser(cnName)
	if !exists {
		return nil
	}

	return []*ldap.Entry{{
		DN: baseDN,
		Attributes: []*ldap.EntryAttribute{
			{Name: "cn", Values: []string{user.CN}},
			{Name: "objectClass", Values: []string{"organizationalPerson"}},
			{Name: "mail", Values: []string{user.Email}},
			{Name: "telephoneNumber", Values: []string{user.Phone}},
		},
	}}
}

func main() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	log.Println("Starting LDAP server...")

	// Create server instance
	s := ldap.NewServer()

	// Register handlers using the correct method names
	s.BindFunc("", &simpleBindHandler{})
	s.SearchFunc("", &simpleSearchHandler{})

	// Start server
	listener, err := net.Listen("tcp", ":389")
	if err != nil {
		log.Fatalf("Listen error: %v", err)
	}
	defer listener.Close()

	// Graceful shutdown
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-quit
		log.Println("Shutting down server...")
		listener.Close()
	}()

	log.Println("Server ready on :389")
	if err := s.Serve(listener); err != nil {
		log.Fatalf("Server error: %v", err)
	}
}
